{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm, trange\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_file(input_file, output_file):\n",
    "    \"\"\"\n",
    "    The function process the data files for Delete & Generate and convert\n",
    "    it for the Delete, Retrieve and Generate training by separating the content\n",
    "    and attributes. It includes all the attribure words.\n",
    "    \n",
    "    Input_file: string : Path of the input file\n",
    "    Output_file: string : Path of the output file \n",
    "    \"\"\"\n",
    "    \n",
    "    with open(input_file) as fp:\n",
    "        data = fp.read().splitlines()\n",
    "    with open (output_file,\"w\") as out_fp:\n",
    "        for x in tqdm(data):\n",
    "            temp = x.split(\"<START>\")\n",
    "            con = temp[0].replace(\"<POS>\",\"\").replace(\"<NEG>\",\"\").replace(\"<CON_START>\",\"\")\n",
    "            sen = temp[1].replace(\"<END>\",\"\")\n",
    "            lt1 = con.split()\n",
    "            lt2 = sen.split()\n",
    "            att_tokens = [z for z in lt2 if z not in lt1]\n",
    "            if len(att_tokens) > max_atts:\n",
    "                max_atts = len(att_tokens)\n",
    "            att_words = \" \".join(att_tokens)\n",
    "            out_str = \"<ATTR_WORDS> \" + att_words + \" <CON_START> \" + con.strip() + \" <START> \" + sen.strip() + \" <END>\" + \"\\n\"\n",
    "            out_fp.write(out_str)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_file_v1(input_file, output_file):\n",
    "    \"\"\"\n",
    "    The function process the data files for Delete & Generate and convert\n",
    "    it for the Delete, Retrieve and Generate training by separating the content\n",
    "    and attributes. It randomly picks 70% of the attributes only to make the generation\n",
    "    more realistic instead of just filling the blanks, which helps while generating\n",
    "    sentences for test cases.\n",
    "    \n",
    "    Input_file: string : Path of the input file\n",
    "    Output_file: string : Path of the output file \n",
    "    \"\"\"\n",
    "    with open(input_file) as fp:\n",
    "        data = fp.read().splitlines()\n",
    "    with open (output_file,\"w\") as out_fp:\n",
    "        for x in tqdm(data):\n",
    "            temp = x.split(\"<START>\")\n",
    "            con = temp[0].replace(\"<POS>\",\"\").replace(\"<NEG>\",\"\").replace(\"<CON_START>\",\"\")\n",
    "            sen = temp[1].replace(\"<END>\",\"\")\n",
    "            lt1 = con.split()\n",
    "            lt2 = sen.split()\n",
    "            att_words = [z for z in lt2 if z not in lt1]\n",
    "            #print(att_words)\n",
    "            \n",
    "            if len(att_words) > 2:\n",
    "                indx = np.array(list(range(len(att_words))))\n",
    "                np.random.shuffle(indx)\n",
    "                att_words = \" \".join([att_words[indx[k]] for k in range(int(0.7 * len(att_words)))])\n",
    "            else: # If attributes less than 2 then keep all the attributes\n",
    "                att_words = \" \".join(att_words)\n",
    "            #print(att_words)\n",
    "            out_str = \"<ATTR_WORDS> \" + att_words + \" <CON_START> \" + con.strip() + \" <START> \" + sen.strip() + \" <END>\" + \"\\n\"\n",
    "            out_fp.write(out_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "process_file(\"./processed_files_with_bert_with_best_head/sentiment_train.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_train_all_attrs.txt\")\n",
    "process_file(\"./processed_files_with_bert_with_best_head/sentiment_train_1.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_train_1_all_attrs.txt\")\n",
    "process_file(\"./processed_files_with_bert_with_best_head/sentiment_train_0.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_train_0_all_attrs.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "process_file(\"./processed_files_with_bert_with_best_head/sentiment_test.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_test_all_attrs.txt\")\n",
    "process_file(\"./processed_files_with_bert_with_best_head/sentiment_test_1.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_test_1_all_attrs.txt\")\n",
    "process_file(\"./processed_files_with_bert_with_best_head/sentiment_test_0.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_test_0_all_attrs.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "process_file(\"./processed_files_with_bert_with_best_head/sentiment_dev.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_dev_all_attrs.txt\")\n",
    "process_file(\"./processed_files_with_bert_with_best_head/sentiment_dev_0.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_dev_1_all_attrs.txt\")\n",
    "process_file(\"./processed_files_with_bert_with_best_head/sentiment_dev_1.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_dev_0_all_attrs.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "process_file_v1(\"./processed_files_with_bert_with_best_head/sentiment_train.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_train.txt\")\n",
    "process_file_v1(\"./processed_files_with_bert_with_best_head/sentiment_train_1.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_train_1.txt\")\n",
    "process_file_v1(\"./processed_files_with_bert_with_best_head/sentiment_train_0.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_train_0.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "process_file_v1(\"./processed_files_with_bert_with_best_head/sentiment_test.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_test.txt\")\n",
    "process_file_v1(\"./processed_files_with_bert_with_best_head/sentiment_test_1.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_test_1.txt\")\n",
    "process_file_v1(\"./processed_files_with_bert_with_best_head/sentiment_test_0.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_test_0.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "process_file_v1(\"./processed_files_with_bert_with_best_head/sentiment_dev.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_dev.txt\")\n",
    "process_file_v1(\"./processed_files_with_bert_with_best_head/sentiment_dev_0.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_dev_1.txt\")\n",
    "process_file_v1(\"./processed_files_with_bert_with_best_head/sentiment_dev_1.txt\",\"./processed_files_with_bert_with_best_head/delete_retrieve_edit_model/sentiment_dev_0.txt\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
